//
// Created by root on 9/12/25.
//

#include "utils.h"
#include <vector>
#include <string>
#include <algorithm>
#include <cctype>
#include "opencv2/opencv.hpp"
#include "TopsInference/TopsInferRuntime.h"

namespace tops_utils {
    bool av_frame_opencv_save(int buffer_type, AVFrame* frame, const std::string& outputPath) {
    if (!frame) {
        std::cerr << "无效的AVFrame" << std::endl;
        return false;
    }

    const bool deviceBuffer = (buffer_type == 1);
    const int width = frame->width;
    const int height = frame->height;
    const int pix_fmt = frame->format;

    // 准备源数据指针与步长。如果在设备端，需要先同步到主机端
    uint8_t* srcData[4] = {nullptr, nullptr, nullptr, nullptr};
    int srcLinesize[4] = {0, 0, 0, 0};
    std::vector<std::vector<uint8_t>> hostPlanes; // 用于承载DtoH后的平面数据

    auto copy_plane_dtoh = [&](int planeIndex, int planeHeight) -> bool {
        if (!frame->data[planeIndex]) { srcData[planeIndex] = nullptr; srcLinesize[planeIndex] = 0; return true; }
        int bytesPerLine = frame->linesize[planeIndex];
        if (bytesPerLine <= 0 || planeHeight <= 0) { srcData[planeIndex] = nullptr; srcLinesize[planeIndex] = 0; return true; }
        int planeSize = bytesPerLine * planeHeight;
        hostPlanes.emplace_back(planeSize);
        auto &buf = hostPlanes.back();
        if (topsMemcpyDtoH(buf.data(), frame->data[planeIndex], planeSize) != topsSuccess) {
            std::cerr << "topsMemcpyDtoH 失败, plane=" << planeIndex << std::endl;
            return false;
        }
        srcData[planeIndex] = buf.data();
        srcLinesize[planeIndex] = bytesPerLine;
        return true;
    };

    if (deviceBuffer) {
        // 仅在 buffer_type==1 时，从设备端同步
        if (pix_fmt == AV_PIX_FMT_YUV420P || pix_fmt == AV_PIX_FMT_YUVJ420P) {
            if (!copy_plane_dtoh(0, height)) return false;
            if (!copy_plane_dtoh(1, height / 2)) return false;
            if (!copy_plane_dtoh(2, height / 2)) return false;
        } else if (pix_fmt == AV_PIX_FMT_BGR24 || pix_fmt == AV_PIX_FMT_RGB24) {
            if (!copy_plane_dtoh(0, height)) return false;
        } else {
            // 默认按单平面处理
            if (!copy_plane_dtoh(0, height)) return false;
        }
    } else {
        // 主机端内存，直接引用
        for (int i = 0; i < 4; ++i) { srcData[i] = frame->data[i]; srcLinesize[i] = frame->linesize[i]; }
    }

    cv::Mat mat; // 输出BGR图像

    if (pix_fmt == AV_PIX_FMT_BGR24) {
        // 若步长等于width*3，可直接clone，否则逐行拷贝
        cv::Mat wrapped(height, width, CV_8UC3, srcData[0], srcLinesize[0]);
        if (wrapped.step == (size_t)width * 3) {
            mat = wrapped.clone();
        } else {
            mat.create(height, width, CV_8UC3);
            for (int y = 0; y < height; ++y) {
                std::memcpy(mat.ptr(y), srcData[0] + y * srcLinesize[0], (size_t)width * 3);
            }
        }
    } else if (pix_fmt == AV_PIX_FMT_RGB24) {
        // 先wrap再转换到BGR
        cv::Mat wrapped(height, width, CV_8UC3, srcData[0], srcLinesize[0]);
        cv::cvtColor(wrapped, mat, cv::COLOR_RGB2BGR);
    } else if (pix_fmt == AV_PIX_FMT_YUV420P || pix_fmt == AV_PIX_FMT_YUVJ420P) {
        SwsContext* swsCtx = sws_getContext(
            width, height, (AVPixelFormat)pix_fmt,
            width, height, AV_PIX_FMT_BGR24,
            SWS_BILINEAR, nullptr, nullptr, nullptr
        );
        if (!swsCtx) {
            std::cerr << "无法创建SWS上下文" << std::endl;
            return false;
        }
        mat.create(height, width, CV_8UC3);
        uint8_t* destData[4] = { mat.data, nullptr, nullptr, nullptr };
        int destLinesize[4] = { static_cast<int>(mat.step), 0, 0, 0 };
        sws_scale(swsCtx, srcData, srcLinesize, 0, height, destData, destLinesize);
        sws_freeContext(swsCtx);
    } else {
        std::cerr << "不支持的像素格式: " << pix_fmt << std::endl;
        return false;
    }

    // 保存为JPG
    if (!cv::imwrite(outputPath, mat)) {
        std::cerr << "无法保存图片到 " << outputPath << std::endl;
        return false;
    }

    std::cout << "图片已成功保存到 " << outputPath << std::endl;
    return true;
}

    bool memory_opencv_save(bool is_gpu,
                            void* data,
                            size_t sizeBytes,
                            std::vector<int64_t> shape,
                            TopsInference::DataType dtype,
                            const std::string& outputPath) {
        if (data == nullptr) {
            std::cerr << "memory_opencv_save: null data" << std::endl;
            return false;
        }
        if (shape.empty()) {
            std::cerr << "memory_opencv_save: empty shape" << std::endl;
            return false;
        }

        // Interpret shape: 3D -> CHW, 2D -> HW (single channel)
        int C = 1, H = 0, W = 0;
        if (shape.size() == 3) {
            C = static_cast<int>(shape[0]);
            H = static_cast<int>(shape[1]);
            W = static_cast<int>(shape[2]);
        } else if (shape.size() == 2) {
            C = 1;
            H = static_cast<int>(shape[0]);
            W = static_cast<int>(shape[1]);
        } else {
            std::cerr << "memory_opencv_save: unsupported shape dims=" << shape.size() << std::endl;
            return false;
        }
        if (C <= 0 || H <= 0 || W <= 0) {
            std::cerr << "memory_opencv_save: invalid shape values CxHxW=" << C << "x" << H << "x" << W << std::endl;
            return false;
        }

        const size_t elemCount = static_cast<size_t>(C) * static_cast<size_t>(H) * static_cast<size_t>(W);
        const size_t bytesIfFloat = elemCount * sizeof(float);
        const size_t bytesIfU8 = elemCount * sizeof(uint8_t);

        // Optional sanity check
        switch (dtype) {
            case TopsInference::DataType::TIF_FP32:
                if (sizeBytes != bytesIfFloat) {
                    std::cerr << "memory_opencv_save: warning FP32 size mismatch, got=" << sizeBytes
                              << ", expected=" << bytesIfFloat << std::endl;
                }
                break;
            case TopsInference::DataType::TIF_UINT8:
                if (sizeBytes != bytesIfU8) {
                    std::cerr << "memory_opencv_save: warning U8 size mismatch, got=" << sizeBytes
                              << ", expected=" << bytesIfU8 << std::endl;
                }
                break;
            default:
                // Fallback paths below will try best-effort based on dtype category
                break;
        }

        // Pull to host if needed
        std::vector<unsigned char> hostBytes;
        const void* hostPtr = nullptr;
        if (is_gpu) {
            hostBytes.resize(sizeBytes);
            if (topsMemcpyDtoH(hostBytes.data(), data, sizeBytes) != topsSuccess) {
                std::cerr << "memory_opencv_save: D2H failed" << std::endl;
                return false;
            }
            hostPtr = hostBytes.data();
        } else {
            hostPtr = data;
        }

        cv::Mat image8u; // final to save
        if (dtype == TopsInference::DataType::TIF_FP32) {
            const float* fptr = reinterpret_cast<const float*>(hostPtr);
            if (C == 1) {
                cv::Mat fimg(H, W, CV_32FC1, const_cast<float*>(fptr));
                // Scale selection: if in [0,1] -> scale 255, else if in [0,255] -> scale 1, else normalize
                double minv, maxv; cv::minMaxLoc(fimg, &minv, &maxv);
                double alpha = 1.0, beta = 0.0;
                if (maxv <= 1.0 + 1e-6 && minv >= 0.0) { alpha = 255.0; }
                else if (maxv <= 255.0 && minv >= 0.0) { alpha = 1.0; }
                else if (maxv > minv) { alpha = 255.0 / (maxv - minv); beta = -minv * alpha; }
                cv::Mat tmp8u; fimg.convertTo(tmp8u, CV_8UC1, alpha, beta);
                image8u = tmp8u;
            } else if (C == 3 || C == 4) {
                size_t planeSize = static_cast<size_t>(H) * static_cast<size_t>(W);
                std::vector<cv::Mat> planes; planes.reserve(C);
                for (int c = 0; c < C; ++c) {
                    float* p = const_cast<float*>(fptr) + c * planeSize;
                    planes.emplace_back(H, W, CV_32FC1, p);
                }
                cv::Mat fimg; cv::merge(planes, fimg); // HWC float, assumed RGB(A)
                // Choose scale similarly
                double minv = 0.0, maxv = 0.0;
                for (int c = 0; c < C; ++c) {
                    double cmin, cmax; cv::minMaxLoc(planes[c], &cmin, &cmax);
                    if (c == 0 || cmin < minv) minv = cmin;
                    if (c == 0 || cmax > maxv) maxv = cmax;
                }
                double alpha = 1.0, beta = 0.0;
                if (maxv <= 1.0 + 1e-6 && minv >= 0.0) { alpha = 255.0; }
                else if (maxv <= 255.0 && minv >= 0.0) { alpha = 1.0; }
                else if (maxv > minv) { alpha = 255.0 / (maxv - minv); beta = -minv * alpha; }
                int cvType = (C == 3) ? CV_8UC3 : CV_8UC4;
                cv::Mat tmp8u; fimg.convertTo(tmp8u, cvType, alpha, beta);
                // Convert RGB(A) -> BGR(A) for saving consistent with OpenCV expectations
                if (C == 3) {
                    cv::cvtColor(tmp8u, image8u, cv::COLOR_RGB2BGR);
                } else {
                    cv::cvtColor(tmp8u, image8u, cv::COLOR_RGBA2BGRA);
                }
            } else {
                std::cerr << "memory_opencv_save: unsupported channels for FP32: C=" << C << std::endl;
                return false;
            }
        } else if (dtype == TopsInference::DataType::TIF_UINT8) { // U8 path, interpreted as CHW RGB or HW
            const uint8_t* uptr = reinterpret_cast<const uint8_t*>(hostPtr);
            if (C == 1) {
                image8u = cv::Mat(H, W, CV_8UC1, const_cast<uint8_t*>(uptr)).clone();
            } else if (C == 3 || C == 4) {
                size_t planeSize = static_cast<size_t>(H) * static_cast<size_t>(W);
                std::vector<cv::Mat> planes; planes.reserve(C);
                for (int c = 0; c < C; ++c) {
                    uint8_t* p = const_cast<uint8_t*>(uptr) + c * planeSize;
                    planes.emplace_back(H, W, CV_8UC1, p);
                }
                cv::Mat rgbOrRgba; cv::merge(planes, rgbOrRgba);
                if (C == 3) {
                    cv::cvtColor(rgbOrRgba, image8u, cv::COLOR_RGB2BGR);
                } else {
                    cv::cvtColor(rgbOrRgba, image8u, cv::COLOR_RGBA2BGRA);
                }
            } else {
                std::cerr << "memory_opencv_save: unsupported channels for U8: C=" << C << std::endl;
                return false;
            }
        } else {
            std::cerr << "memory_opencv_save: unsupported dtype " << static_cast<int>(dtype) << std::endl;
            return false;
        }

        if (image8u.empty()) {
            std::cerr << "memory_opencv_save: empty image after conversion" << std::endl;
            return false;
        }

        std::vector<int> params;
        // reasonable defaults
        std::string lower = outputPath; std::transform(lower.begin(), lower.end(), lower.begin(), [](unsigned char c){ return (char)std::tolower(c); });
        if (lower.size() >= 4) {
            if (lower.rfind(".jpg") != std::string::npos || lower.rfind(".jpeg") != std::string::npos) {
                params = {cv::IMWRITE_JPEG_QUALITY, 95};
            } else if (lower.rfind(".png") != std::string::npos) {
                params = {cv::IMWRITE_PNG_COMPRESSION, 3};
            }
        }

        if (!cv::imwrite(outputPath, image8u, params)) {
            std::cerr << "memory_opencv_save: failed to write image: " << outputPath << std::endl;
            return false;
        }
        return true;
    }
}